package com.company.nuwa.redis.support.limit;

import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;
import org.redisson.api.*;
import org.springframework.util.Assert;

import java.util.concurrent.Callable;

/**
 * 基于Redis实现的限流器(滑动时间窗口){@link RRateLimiter}
 * <p>
 * 亦可参考guava基于JVM实现版本 {@link com.google.common.util.concurrent.RateLimiter}
 *
 * @author Nikola Tesla
 * @version 1.0
 * @date 2020/09/18 13:42
 */
@Slf4j
public class RedisRateLimiter {

    private final RedissonClient redissonClient;

    public RedisRateLimiter(RedissonClient redissonClient) {
        this.redissonClient = redissonClient;
    }


    /**
     * @param limitKey     lock key
     * @param rate         rate
     * @param rateInterval rate time interval (ms)
     * @param callable     callable
     * @return callable value
     */
    @SneakyThrows
    public <V> V acquire(String limitKey, long rate, long rateInterval, Callable<V> callable) {
        Assert.notNull(callable, "Callback object must not be null");

        RRateLimiter rateLimiter = this.getRateLimiter(limitKey, rate, rateInterval);

        rateLimiter.acquire();

        return callable.call();
    }

    /**
     * @param limitKey     lock key
     * @param rate         rate
     * @param rateInterval rate time interval (ms)
     * @param success      success callable
     * @return callable value
     */
    @SneakyThrows
    public <V> V tryAcquire(String limitKey, long rate, long rateInterval, Callable<V> success) {
        Assert.notNull(success, "Callback object must not be null");

        RRateLimiter rateLimiter = this.getRateLimiter(limitKey, rate, rateInterval);

        if (rateLimiter.tryAcquire()) {
            return success.call();
        }
        return null;
    }


    /**
     * @param limitKey     lock key
     * @param rate         rate
     * @param rateInterval rate time interval (ms)
     * @param success      success callable
     * @param fail         failure callable
     * @return callable value
     */
    @SneakyThrows
    public <V> V tryAcquire(String limitKey, long rate, long rateInterval, Callable<V> success, Callable<V> fail) {
        Assert.notNull(success, "Callback object must not be null");

        RRateLimiter rateLimiter = this.getRateLimiter(limitKey, rate, rateInterval);

        if (rateLimiter.tryAcquire()) {
            return success.call();
        } else {
            if (fail != null) {
                return fail.call();
            }
        }

        return null;
    }

    /**
     * 获取 rate limiter
     *
     * @param limitKey
     * @param rate
     * @param rateInterval
     * @return
     */
    private RRateLimiter getRateLimiter(String limitKey, long rate, long rateInterval) {
        RRateLimiter rateLimiter = redissonClient.getRateLimiter(limitKey);
        RateLimiterConfig config = rateLimiter.getConfig();
        // 速率有变更，则重新设置
        if (config.getRate() != rate || config.getRateInterval() != rateInterval || config.getRateType() != RateType.PER_CLIENT) {
            rateLimiter.setRate(RateType.PER_CLIENT, rate, rateInterval, RateIntervalUnit.MILLISECONDS);
        } else {
            rateLimiter.trySetRate(RateType.PER_CLIENT, rate, rateInterval, RateIntervalUnit.MILLISECONDS);
        }
        return rateLimiter;
    }


}
